/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.solr.handler;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.Date;
import java.util.Locale;

import org.apache.solr.common.SolrException;
import org.apache.solr.common.params.CommonParams;
import org.apache.solr.common.params.ModifiableSolrParams;
import org.apache.solr.common.params.SolrParams;
import org.apache.solr.common.util.NamedList;
import org.apache.solr.core.SolrCore;
import org.apache.solr.util.plugin.SolrCoreAware;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.request.SolrRequestHandler;
import org.apache.solr.response.SolrQueryResponse;
import org.apache.solr.schema.DateField;

import org.apache.commons.io.FileUtils;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Ping Request Handler for reporting SolrCore health to a Load Balancer.
 * 
 * <p>
 * This handler is designed to be used as the endpoint for an HTTP Load-Balancer
 * to use when checking the "health" or "up status" of a Solr server.
 * </p>
 * 
 * <p>
 * In it's simplest form, the PingRequestHandler should be configured with some
 * defaults indicating a request that should be executed. If the request
 * succeeds, then the PingRequestHandler will respond back with a simple "OK"
 * status. If the request fails, then the PingRequestHandler will respond back
 * with the corrisponding HTTP Error code. Clients (such as load balancers) can
 * be configured to poll the PingRequestHandler monitoring for these types of
 * responses (or for a simple connection failure) to know if there is a problem
 * with the Solr server.
 * </p>
 * 
 * <pre class="prettyprint">
 * &lt;requestHandler name="/admin/ping" class="solr.PingRequestHandler"&gt;
 *   &lt;lst name="invariants"&gt;
 *     &lt;str name="qt"&gt;/search&lt;/str&gt;&lt;!-- handler to delegate to --&gt;
 *     &lt;str name="q"&gt;some test query&lt;/str&gt;
 *   &lt;/lst&gt;
 * &lt;/requestHandler&gt;
 * </pre>
 * 
 * <p>
 * A more advanced option available, is to configure the handler with a
 * "healthcheckFile" which can be used to enable/disable the PingRequestHandler.
 * </p>
 * 
 * <pre class="prettyprint">
 * &lt;requestHandler name="/admin/ping" class="solr.PingRequestHandler"&gt;
 *   &lt;!-- relative paths are resolved against the data dir --&gt;
 *   &lt;str name="healthcheckFile"&gt;server-enabled.txt&lt;/str&gt;
 *   &lt;lst name="invariants"&gt;
 *     &lt;str name="qt"&gt;/search&lt;/str&gt;&lt;!-- handler to delegate to --&gt;
 *     &lt;str name="q"&gt;some test query&lt;/str&gt;
 *   &lt;/lst&gt;
 * &lt;/requestHandler&gt;
 * </pre>
 * 
 * <ul>
 * <li>If the health check file exists, the handler will execute the delegated
 * query and return status as described above.</li>
 * <li>If the health check file does not exist, the handler will return an HTTP
 * error even if the server is working fine and the delegated query would have
 * succeeded</li>
 * </ul>
 * 
 * <p>
 * This health check file feature can be used as a way to indicate to some Load
 * Balancers that the server should be "removed from rotation" for maintenance,
 * or upgrades, or whatever reason you may wish.
 * </p>
 * 
 * <p>
 * The health check file may be created/deleted by any external system, or the
 * PingRequestHandler itself can be used to create/delete the file by specifying
 * an "action" param in a request:
 * </p>
 * 
 * <ul>
 * <li><code>http://.../ping?action=enable</code> - creates the health check
 * file if it does not already exist</li>
 * <li><code>http://.../ping?action=disable</code> - deletes the health check
 * file if it exists</li>
 * <li><code>http://.../ping?action=status</code> - returns a status code
 * indicating if the healthcheck file exists ("<code>enabled</code>") or not ("
 * <code>disabled<code>")</li>
 * </ul>
 * 
 * @since solr 1.3
 */
public class PingRequestHandler extends RequestHandlerBase implements
		SolrCoreAware {
	public static Logger log = LoggerFactory
			.getLogger(PingRequestHandler.class);

	public static final String HEALTHCHECK_FILE_PARAM = "healthcheckFile";

	protected enum ACTIONS {
		STATUS, ENABLE, DISABLE, PING
	};

	private String healthFileName = null;
	private File healthcheck = null;

	public void init(NamedList args) {
		super.init(args);
		Object tmp = args.get(HEALTHCHECK_FILE_PARAM);
		healthFileName = (null == tmp ? null : tmp.toString());
	}

	public void inform(SolrCore core) {
		if (null != healthFileName) {
			healthcheck = new File(healthFileName);
			if (!healthcheck.isAbsolute()) {
				healthcheck = new File(core.getDataDir(), healthFileName);
				healthcheck = healthcheck.getAbsoluteFile();
			}

			if (!healthcheck.getParentFile().canWrite()) {
				// this is not fatal, users may not care about enable/disable
				// via
				// solr request, file might be touched/deleted by an external
				// system
				log.warn(
						"Directory for configured healthcheck file is not writable by solr, PingRequestHandler will not be able to control enable/disable: {}",
						healthcheck.getParentFile().getAbsolutePath());
			}

		}

	}

	/**
	 * Returns true if the healthcheck flag-file is enabled but does not exist,
	 * otherwise (no file configured, or file configured and exists) returns
	 * false.
	 */
	public boolean isPingDisabled() {
		return (null != healthcheck && !healthcheck.exists());
	}

	@Override
	public void handleRequestBody(SolrQueryRequest req, SolrQueryResponse rsp)
			throws Exception {

		SolrParams params = req.getParams();

		// in this case, we want to default distrib to false so
		// we only ping the single node
		Boolean distrib = params.getBool("distrib");
		if (distrib == null) {
			ModifiableSolrParams mparams = new ModifiableSolrParams(params);
			mparams.set("distrib", false);
			req.setParams(mparams);
		}

		String actionParam = params.get("action");
		ACTIONS action = null;
		if (actionParam == null) {
			action = ACTIONS.PING;
		} else {
			try {
				action = ACTIONS.valueOf(actionParam.toUpperCase(Locale.ROOT));
			} catch (IllegalArgumentException iae) {
				throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
						"Unknown action: " + actionParam);
			}
		}
		switch (action) {
		case PING:
			if (isPingDisabled()) {
				throw new SolrException(
						SolrException.ErrorCode.SERVICE_UNAVAILABLE,
						"Service disabled");
			}
			handlePing(req, rsp);
			break;
		case ENABLE:
			handleEnable(true);
			break;
		case DISABLE:
			handleEnable(false);
			break;
		case STATUS:
			if (healthcheck == null) {
				SolrException e = new SolrException(
						SolrException.ErrorCode.SERVICE_UNAVAILABLE,
						"healthcheck not configured");
				rsp.setException(e);
			} else {
				rsp.add("status", isPingDisabled() ? "disabled" : "enabled");
			}
		}
	}

	protected void handlePing(SolrQueryRequest req, SolrQueryResponse rsp)
			throws Exception {

		SolrParams params = req.getParams();
		SolrCore core = req.getCore();

		// Get the RequestHandler
		String qt = params.get(CommonParams.QT);// optional; you get the default
												// otherwise
		SolrRequestHandler handler = core.getRequestHandler(qt);
		if (handler == null) {
			throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
					"Unknown RequestHandler (qt): " + qt);
		}

		if (handler instanceof PingRequestHandler) {
			throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
					"Cannot execute the PingRequestHandler recursively");
		}

		// Execute the ping query and catch any possible exception
		Throwable ex = null;
		try {
			SolrQueryResponse pingrsp = new SolrQueryResponse();
			core.execute(handler, req, pingrsp);
			ex = pingrsp.getException();
		} catch (Throwable th) {
			ex = th;
		}

		// Send an error or an 'OK' message (response code will be 200)
		if (ex != null) {
			throw new SolrException(SolrException.ErrorCode.SERVER_ERROR,
					"Ping query caused exception: " + ex.getMessage(), ex);
		}
		rsp.add("status", "OK");
	}

	protected void handleEnable(boolean enable) throws SolrException {
		if (healthcheck == null) {
			throw new SolrException(
					SolrException.ErrorCode.SERVICE_UNAVAILABLE,
					"No healthcheck file defined.");
		}
		if (enable) {
			try {
				// write out when the file was created
				FileUtils.write(healthcheck,
						DateField.formatExternal(new Date()), "UTF-8");
			} catch (IOException e) {
				throw new SolrException(SolrException.ErrorCode.SERVER_ERROR,
						"Unable to write healthcheck flag file", e);
			}
		} else {
			if (healthcheck.exists() && !healthcheck.delete()) {
				throw new SolrException(SolrException.ErrorCode.NOT_FOUND,
						"Did not successfully delete healthcheck file: "
								+ healthcheck.getAbsolutePath());
			}
		}
	}

	// ////////////////////// SolrInfoMBeans methods //////////////////////

	@Override
	public String getDescription() {
		return "Reports application health to a load-balancer";
	}

	@Override
	public String getSource() {
		return "$URL: https://svn.apache.org/repos/asf/lucene/dev/branches/lucene_solr_4_0/solr/core/src/java/org/apache/solr/handler/PingRequestHandler.java $";
	}
}
